use hmac::{Hmac, Mac};
use pbkdf2::pbkdf2;
use rand::random;
use sha2::{Digest, Sha256};

type HmacSha256 = Hmac<Sha256>;

fn bytes_32_from_slice(slice: &[u8]) -> [u8; 32] {
    let mut out = [0u8; 32];
    out.copy_from_slice(slice);
    out
}

fn hmac_sha256(key: &[u8], input: &[u8]) -> [u8; 32] {
    let mut mac = HmacSha256::new_from_slice(key).expect("client key hmac init error");
    mac.update(input);
    bytes_32_from_slice(&*mac.finalize().into_bytes())
}

pub fn rand_32_bytes() -> [u8; 32] {
    random()
}

pub fn salted_password(password: &[u8], salt: &[u8], iteration: usize) -> [u8; 32] {
    let mut buf = [0u8; 32];
    pbkdf2::<HmacSha256>(password, salt, iteration as u32, &mut buf).expect("PBKDF2 fail");
    buf
}

pub fn client_key(input: &[u8]) -> [u8; 32] {
    hmac_sha256(b"Client Key", input)
}

pub fn _server_key(input: &[u8]) -> [u8; 32] {
    hmac_sha256(b"Server Key", input)
}

pub fn stored_key(input: &[u8]) -> [u8; 32] {
    let mut sha = Sha256::new();
    sha.update(input);
    bytes_32_from_slice(&*sha.finalize())
}

pub fn signature(stored_key: &[u8], authmsg: &[u8]) -> [u8; 32] {
    hmac_sha256(authmsg, stored_key)
}

pub fn _client_proof(password: &[u8], salt: &[u8], iteration: usize, authmsg: &[u8]) -> [u8; 32] {
    let pw = salted_password(password, salt, iteration);
    let mut ck = client_key(&pw);
    let sk = stored_key(&ck);
    let sig = signature(&sk, authmsg);

    for (ck_b, sig_b) in ck.iter_mut().zip(sig.iter().copied()) {
        *ck_b ^= sig_b
    }
    ck
}

pub fn _server_proof(password: &[u8], salt: &[u8], iteration: usize, authmsg: &[u8]) -> [u8; 32] {
    let pw = salted_password(password, salt, iteration);
    let sk = _server_key(&pw);
    signature(&sk, authmsg)
}
